#include "Plot2D.h"

//<Analysis::Initialize>
Plot2D::Plot2D() {
    string channels[Nch]={"eeee", "eemm", "mmee", "mmmm", "2e2m", "incl"};
    for(int i=0; i<Nch; i++) { ch.push_back(channels[i]); }
    froot = new TFile("result.root", "RECREATE");
    for( int i=0; i<Nch; i++ ) {
        // clear counter
        Ncount_data[i] = 0; Ncount_mc[i] = 0;
        Ncount_Cut_data[i] = 0; Ncount_Cut_mc[i] = 0;

        string hname;
        // ********** plots without m4l cut
        // plots for data
        hname = "MZZMZ1_data_" + ch[i];
    	h2DMZZMZ1_data[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
        hname = "MZZMZ2_data_" + ch[i];
    	h2DMZZMZ2_data[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
        hname = "MZ1MZ2_data_" + ch[i];
    	h2DMZ1MZ2_data[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
        // plots for MC
        hname = "MZZMZ1_mc_" + ch[i];
    	h2DMZZMZ1_mc[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
        hname = "MZZMZ2_mc_" + ch[i];
    	h2DMZZMZ2_mc[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
        hname = "MZ1MZ2_mc_" + ch[i];
    	h2DMZ1MZ2_mc[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
        // ********** plots after m4l cut, 118-130
        // plots for data
        hname = "MZZMZ1_Cut_data_" + ch[i];
    	h2DMZZMZ1_Cut_data[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
        hname = "MZZMZ2_Cut_data_" + ch[i];
    	h2DMZZMZ2_Cut_data[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
        hname = "MZ1MZ2_Cut_data_" + ch[i];
    	h2DMZ1MZ2_Cut_data[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
        // plots for MC
        hname = "MZZMZ1_Cut_mc_" + ch[i];
    	h2DMZZMZ1_Cut_mc[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
        hname = "MZZMZ2_Cut_mc_" + ch[i];
    	h2DMZZMZ2_Cut_mc[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
        hname = "MZ1MZ2_Cut_mc_" + ch[i];
    	h2DMZ1MZ2_Cut_mc[i] = new TH2F(hname.c_str(), hname.c_str(), 2000, 0, 200, 2000, 0, 200);
    }
}

//<Analysis::LoadEntry>
bool Plot2D::LoadEntry(long entry) {
    fChain->GetEntry(entry);
}

//<Analysis::ProcessEntry>
bool Plot2D::ProcessEntry(long entry, string stream, int nch) {
    if(stream=="data") {
        cout << "stream/ch/m4l: " << stream << "/" << nch << "/" << m4l << std::endl;

        // counting
        Ncount_data[nch] += weight;
        Ncount_data[5] += weight;
        if(nch==1||nch==2) Ncount_data[4] += weight;

        // push back to vector
        vec_MZ1[nch].push_back(m1); vec_MZ2[nch].push_back(m2); vec_MZZ[nch].push_back(m4l);
        vec_MZ1[5].push_back(m1); vec_MZ2[5].push_back(m2); vec_MZZ[5].push_back(m4l);
        if(nch==1||nch==2) {
        	vec_MZ1[4].push_back(m1); vec_MZ2[4].push_back(m2); vec_MZZ[4].push_back(m4l);
        }
 
        // fill histos
        h2DMZZMZ1_data[nch]->Fill(m4l, m1, weight);
        h2DMZZMZ1_data[5]->Fill(m4l, m1, weight); //inclusive
        if(nch==1||nch==2) h2DMZZMZ1_data[4]->Fill(m4l, m1, weight); //2e2m

        h2DMZZMZ2_data[nch]->Fill(m4l, m2, weight);
        h2DMZZMZ2_data[5]->Fill(m4l, m2, weight); //inclusive
        if(nch==1||nch==2) h2DMZZMZ2_data[4]->Fill(m4l, m2, weight); //2e2m

        h2DMZ1MZ2_data[nch]->Fill(m1, m2, weight);
        h2DMZ1MZ2_data[5]->Fill(m1, m2, weight); //inclusive
        if(nch==1||nch==2) h2DMZ1MZ2_data[4]->Fill(m1, m2, weight); //2e2m
        // m4l cut
        if(m4l>118 && m4l<130) {
        	Ncount_Cut_data[nch] += weight;
        	Ncount_Cut_data[5] += weight;
        	if(nch==1||nch==2) Ncount_Cut_data[4] += weight;

        	// push back to vector
        	vec_Cut_MZ1[nch].push_back(m1); vec_Cut_MZ2[nch].push_back(m2); vec_Cut_MZZ[nch].push_back(m4l);
        	vec_Cut_MZ1[5].push_back(m1); vec_Cut_MZ2[5].push_back(m2); vec_Cut_MZZ[5].push_back(m4l);
        	if(nch==1||nch==2) {
        		vec_Cut_MZ1[4].push_back(m1); vec_Cut_MZ2[4].push_back(m2); vec_Cut_MZZ[4].push_back(m4l);
        	}
 
        	h2DMZZMZ1_Cut_data[nch]->Fill(m4l, m1, weight);
        	h2DMZZMZ1_Cut_data[5]->Fill(m4l, m1, weight); //inclusive
        	if(nch==1||nch==2) h2DMZZMZ1_Cut_data[4]->Fill(m4l, m1, weight); //2e2m

        	h2DMZZMZ2_Cut_data[nch]->Fill(m4l, m2, weight);
        	h2DMZZMZ2_Cut_data[5]->Fill(m4l, m2, weight); //inclusive
        	if(nch==1||nch==2) h2DMZZMZ2_Cut_data[4]->Fill(m4l, m2, weight); //2e2m

        	h2DMZ1MZ2_Cut_data[nch]->Fill(m1, m2, weight);
        	h2DMZ1MZ2_Cut_data[5]->Fill(m1, m2, weight); //inclusive
        	if(nch==1||nch==2) h2DMZ1MZ2_Cut_data[4]->Fill(m1, m2, weight); //2e2m
        }
    }
    if(stream=="mc") {
        cout << "stream/ch/m4l: " << stream << "/" << nch << "/" << m4l << std::endl;

        Ncount_mc[nch] += weight;
        Ncount_mc[5] += weight;
        if(nch==1||nch==2) Ncount_mc[4] += weight;

        h2DMZZMZ1_mc[nch]->Fill(m4l, m1, weight);
        h2DMZZMZ1_mc[5]->Fill(m4l, m1, weight); //inclusive
        if(nch==1||nch==2) h2DMZZMZ1_mc[4]->Fill(m4l, m1, weight); //2e2m

        h2DMZZMZ2_mc[nch]->Fill(m4l, m2, weight);
        h2DMZZMZ2_mc[5]->Fill(m4l, m2, weight); //inclusive
        if(nch==1||nch==2) h2DMZZMZ2_mc[4]->Fill(m4l, m2, weight); //2e2m

        h2DMZ1MZ2_mc[nch]->Fill(m1, m2, weight);
        h2DMZ1MZ2_mc[5]->Fill(m1, m2, weight); //inclusive
        if(nch==1||nch==2) h2DMZ1MZ2_mc[4]->Fill(m1, m2, weight); //2e2m
        if(m4l>118 && m4l<130) {
        	Ncount_Cut_mc[nch] += weight;
        	Ncount_Cut_mc[5] += weight;
        	if(nch==1||nch==2) Ncount_Cut_mc[4] += weight;

        	h2DMZZMZ1_Cut_mc[nch]->Fill(m4l, m1, weight);
        	h2DMZZMZ1_Cut_mc[5]->Fill(m4l, m1, weight); //inclusive
        	if(nch==1||nch==2) h2DMZZMZ1_Cut_mc[4]->Fill(m4l, m1, weight); //2e2m

        	h2DMZZMZ2_Cut_mc[nch]->Fill(m4l, m2, weight);
        	h2DMZZMZ2_Cut_mc[5]->Fill(m4l, m2, weight); //inclusive
        	if(nch==1||nch==2) h2DMZZMZ2_Cut_mc[4]->Fill(m4l, m2, weight); //2e2m

        	h2DMZ1MZ2_Cut_mc[nch]->Fill(m1, m2, weight);
        	h2DMZ1MZ2_Cut_mc[5]->Fill(m1, m2, weight); //inclusive
        	if(nch==1||nch==2) h2DMZ1MZ2_Cut_mc[4]->Fill(m1, m2, weight); //2e2m
        }
    }
}

//<Analysis::SlaveTerminate>
void Plot2D::Terminate() {
    // output
    FILE *fout = fopen("log_eff_out.txt", "w");
    fprintf(fout, "data: \n");
    for( int i=0; i<Nch; i++ ) {
        fprintf(fout, "%s: %.2f\n", ch[i].c_str(), Ncount_data[i]);
    }
    fprintf(fout, "data with 118 < m4l < 130: \n");
    for( int i=0; i<Nch; i++ ) {
        fprintf(fout, "%s: %.2f\n", ch[i].c_str(), Ncount_Cut_data[i]);
    }
    fprintf(fout, "mc: \n");
    for( int i=0; i<Nch; i++ ) {
        fprintf(fout, "%s: %.2f\n", ch[i].c_str(), Ncount_mc[i]);
    }
    fprintf(fout, "mc with 118 < m4l < 130: \n");
    for( int i=0; i<Nch; i++ ) {
        fprintf(fout, "%s: %.2f\n", ch[i].c_str(), Ncount_Cut_mc[i]);
    }
    fprintf(fout, "\n mass for data: \n");
    for( int i=0; i<Nch; i++ ) {
        fprintf(fout, "%s:\n", ch[i].c_str());
        fprintf(fout, "MZZ: ");
        for( int j=0; j<(int)vec_Cut_MZZ[i].size(); j++ ) {
        	fprintf(fout, "%10.2f,", vec_Cut_MZZ[i][j]);
     	}   
        fprintf(fout, "\n");
        fprintf(fout, "MZ1: ");
        for( int j=0; j<(int)vec_Cut_MZ1[i].size(); j++ ) {
        	fprintf(fout, "%10.2f,", vec_Cut_MZ1[i][j]);
     	}   
        fprintf(fout, "\n");
        fprintf(fout, "MZ2: ");
        for( int j=0; j<(int)vec_Cut_MZ2[i].size(); j++ ) {
        	fprintf(fout, "%10.2f,", vec_Cut_MZ2[i][j]);
     	}   
        fprintf(fout, "\n");
    }
     

    
    froot->Write();
    // plot
}
